import sqlite3
from typing import TypedDict
from langgraph.graph import StateGraph,START
from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.sqlite import SqliteSaver

class State(TypedDict):
    topic:str
    joke:str

def refine_topic(state:State):
    return {"topic":state["topic"] + " and cats"}

def generate_joke(state:State):
    return {"joke":f"This is a joke about {state['topic']}"}
conn = sqlite3.connect("langgraph.db",check_same_thread=False)
checkpointer = SqliteSaver(conn)
config = {"configurable":{"thread_id":"1"}}
graph = (StateGraph(State)
         .add_node(refine_topic).add_node(generate_joke)
    .add_edge(START,"refine_topic")
    .add_edge("refine_topic","generate_joke")
    .compile(checkpointer=checkpointer))


# checkpoint = MemorySaver()
# graph.checkpointer(checkpointer)

for chunk in graph.stream({"topic":"ice_cream"},stream_mode="values",config={"configurable": {"thread_id": "1808"}}):
    print(chunk)

